import numpy as np
from scipy.signal import argrelextrema
from scipy.optimize import curve_fit
from matplotlib import animation, pyplot as plt
from IPython import display
from src.simulation import LatticeBoltzmann
size_x = 80
size_y = 20
n_steps = 8000
omega = 1.3
epsilon = 0.01
rho0 = 0.5
# initialize density
density_x = rho0 + epsilon * np.sin(2*np.pi*np.arange(size_x)/size_x)
density = np.tile(density_x, (size_y, 1))
# create lattice
lattice = LatticeBoltzmann(size_x, size_y, omega=omega, init_density=density)
sine = np.empty((n_steps+1, lattice.X))
sine[0] = lattice.density[0]
for i in range(n_steps):
lattice.step()
# collect the sinusoidal density profile of the first row in each step
sine[i+1] = lattice.density[0]
# visualize
lattice.plot.display_animation()